package com.lw.leetcode.tree.b;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * Created with IntelliJ IDEA.
 * 1519. 子树中标签相同的节点数
 *
 * @author liw
 * @version 1.0
 * @date 2022/4/20 14:18
 */
public class CountSubTrees {

    private int[] counts;

    public int[] countSubTrees(int n, int[][] edges, String labels) {
        Map<Integer, Node> map = new HashMap<>();
        counts = new int[n];
        for (int[] edge : edges) {
            int a = edge[0];
            int b = edge[1];
            Node bn = map.computeIfAbsent(b, v -> new Node(b, labels.charAt(b)));
            Node an = map.computeIfAbsent(a, v -> new Node(a, labels.charAt(a)));
            an.list.add(bn);
            bn.list.add(an);
        }
        change(map.get(0), -1);
        find(map.get(0));
        return counts;
    }

    private void change(Node node, int f) {
        List<Node> list = node.list;
        for (int size = list.size() - 1; size >= 0; size--) {
            if (list.get(size).val == f) {
                node.list.remove(size);
            } else {
                change(list.get(size), node.val);
            }
        }
    }

    private int[] find(Node node) {
        int[] arr = null;
        if (node.list.isEmpty()) {
            arr = new int[26];
        } else {
            for (Node no : node.list) {
                int[] ints = find(no);
                if (arr == null) {
                    arr = ints;
                } else {
                    for (int i = 0; i < 26; i++) {
                        arr[i] += ints[i];
                    }
                }
            }
        }
        arr[node.c - 'a']++;
        counts[node.val] = arr[node.c - 'a'];
        return arr;
    }

    class Node {
        public Node(int val, char c) {
            this.val = val;
            this.c = c;
        }

        private int val;
        private char c;
        private List<Node> list = new ArrayList<>();
    }

}
